import pandas as pd
import numpy as np

from sklearn.ensemble import RandomForestRegressor
from sklearn.model_selection import train_test_split

# https://chat.deepseek.com/a/chat/s/6fb7bd95-582e-4fbc-a7b6-5b92b20d1d97
# 历史数据示例结构
history_data = pd.DataFrame({
    'item_id': [1,1,1,1,1,1,1, 2,2,2,2,2,2,2], # 商品id
    'day_num': [0,1,2,3,4,5,6, 0,1,2,3,4,5,6],  # 0表示上新日
    'sale': [200,180,150,120,100,80,60, 90,70,65,60,50,40,30], #销量
    'pre_add': [1500,1500,1500,1500,1500,1500,1500, 800,800,800,800,800,800,800] # 上新前累计加购数？ 还是加购数
})

# 新款实时数据
new_item_data = {
    'item_id': 3,  # 商品id
    'pre_add': 1700,  # 累计加购数
    'sales': {0: 280}  # 每日销量字典，key为day_num
}

# pre_add 累计加购数
# sales 日销
# sale_decay 销量衰减
# cum_sale 累计销量
#
# sale_ratio 累计销量/加购数

def create_features(df, current_day):
    """创建滚动预测特征"""
    # 基础特征
    df['days_since_launch'] = df['day_num'] - current_day

    # 时序特征  //先分组 然后对分组的数据进行分别操作
    df['sale_lag1'] = df.groupby('item_id')['sale'].shift(1)
    df['sale_decay'] = df['sale'] / (df['sale_lag1'] + 1e-5)  # 避免除零

    # 累计特征
    df['cum_sale'] = df.groupby('item_id')['sale'].cumsum()
    df['sale_ratio'] = df['cum_sale'] / df['pre_add']

    # # 在create_features函数中添加：
    # df['sale_ma3'] = df.groupby('item_id')['sale'].transform(
    #     lambda x: x.rolling(3, min_periods=1).mean())
    #
    # df['sale_trend'] = df.groupby('item_id')['sale'].transform(
    #     lambda x: x.diff().rolling(3, min_periods=1).mean())
    #
    # df['day_of_week'] = df['day_num'] % 7

    # 时间特征
    # df['is_weekend'] = ((df['day_num'] % 7) >= 5).astype(int)  # 假设day0是周一

    # # 未来预测目标
    # df['future_7d_sale'] = df.groupby('item_id')['sale'].transform(
    #     lambda x: x.shift(-1).rolling(7, min_periods=1).sum())

    # return df.dropna(subset=['future_7d_sale'])

    return df

def train_base_model(history_data):
    """使用历史数据训练基础模型"""
    # 创建特征
    train_df = create_features(history_data.copy(), current_day=0)

    # 特征和目标
    features = ['pre_add', 'sale_lag1', 'sale_decay',  'sale_ratio',
                # 'is_weekend',
                'days_since_launch']
    X = train_df[features]
    y = train_df['cum_sale']

    # 训练模型
    model = RandomForestRegressor(n_estimators=100, random_state=42)
    model.fit(X, y)

    return model, features


def daily_rolling_forecast(model, features, item_data, current_day):
    """
    每日滚动预测未来7天销量
    model: 训练好的预测模型
    features: 模型使用的特征列表
    item_data: 商品数据（包含历史销售记录）
    current_day: 当前天数（上新日为0）
    """
    # 准备预测数据
    future_days = pd.DataFrame({
        'item_id': [item_data['item_id']] * 7,
        'day_num': range(current_day + 1, current_day + 8),
        'pre_add': [item_data['pre_add']] * 7
    })

    # 添加已知销量
    for day, sale in item_data['sales'].items():
        if day < current_day:
            future_days = future_days.append({
                'item_id': item_data['item_id'],
                'day_num': day,
                'sale': sale,
                'pre_add': item_data['pre_add']
            }, ignore_index=True)

    # 创建特征
    future_days = create_features(future_days.sort_values('day_num'), current_day)

    # 预测未来7天销量
    future_days = future_days[future_days['day_num'] > current_day]
    future_days['pred_sale'] = model.predict(future_days[features])

    return future_days[['day_num', 'pred_sale']]


def update_model(model, history_data, new_data, current_day):
    """使用新增数据更新模型"""
    # 添加新数据到历史数据集
    new_rows = []
    for day, sale in new_data['sales'].items():
        if day <= current_day:  # 只添加已发生的日期
            new_rows.append({
                'item_id': new_data['item_id'],
                'day_num': day,
                'sale': sale,
                'pre_add': new_data['pre_add']
            })

    updated_data = pd.concat([history_data, pd.DataFrame(new_rows)], ignore_index=True)

    # 重新训练模型
    updated_df = create_features(updated_data.copy(), current_day)
    X = updated_df[features]
    y = updated_df['future_7d_sale']

    # 增量学习（或完全重训）
    if hasattr(model, 'partial_fit'):
        model.partial_fit(X, y)
    else:
        model.fit(X, y)  # 对于不支持增量学习的模型，完全重训

    return model, updated_data


# 初始化
base_model, features = train_base_model(history_data)
all_data = history_data.copy()
current_day = 0

# 模拟每日滚动预测
for day in range(0, 14):  # 模拟两周的每日预测
    print(f"\n=== Day {day} ===")

    # 获取当日真实销量（模拟）
    if day > 0:  # 上新日之后才有新数据
        # 模拟生成当日销量（实际应用中从数据库获取）
        # todo
        new_sale = generate_daily_sale(new_item_data, day)
        new_item_data['sales'][day] = new_sale
        # new_sale= new_sale

    # 每日滚动预测
    forecast = daily_rolling_forecast(base_model, features, new_item_data, day)
    print(f"预测 {day + 1} 到 {day + 7} 天销量:")
    print(forecast.set_index('day_num')['pred_sale'])

    # 如果有新数据，更新模型
    if day > 0:
        base_model, all_data = update_model(
            base_model,
            all_data,
            new_item_data,
            day
        )

    # 检查是否收集到新数据
    if day in new_item_data['sales']:
        print(f"当日真实销量: {new_item_data['sales'][day]}")

    # 模拟等待下一天...